1 files changed, 23 insertions, 5 deletions
diff --git a/src/video_core/control/channel_state_cache.h b/src/video_core/control/channel_state_cache.h
index c8298c003..c51040c83 100644
--- a/src/video_core/control/channel_state_cache.h
+++ b/src/video_core/control/channel_state_cache.h
@@ -2,6 +2,7 @@
 
 #include <deque>
 #include <limits>
+#include <mutex>
 #include <unordered_map>
 
 #include "common/common_types.h"
@@ -41,9 +42,10 @@ template <class P>
 class ChannelSetupCaches {
 public:
     /// Operations for seting the channel of execution.
+    virtual ~ChannelSetupCaches();
 
     /// Create channel state.
-    void CreateChannel(Tegra::Control::ChannelState& channel);
+    virtual void CreateChannel(Tegra::Control::ChannelState& channel);
 
     /// Bind a channel for execution.
     void BindToChannel(s32 id);
@@ -51,18 +53,34 @@ public:
     /// Erase channel's state.
     void EraseChannel(s32 id);
 
+    Tegra::MemoryManager* GetFromID(size_t id) const {
+        std::unique_lock<std::mutex> lk(config_mutex);
+        const auto ref = address_spaces.find(id);
+        return ref->second.gpu_memory;
+    }
+
 protected:
     static constexpr size_t UNSET_CHANNEL{std::numeric_limits<size_t>::max()};
 
-    std::deque<P> channel_storage;
-    std::deque<size_t> free_channel_ids;
-    std::unordered_map<s32, size_t> channel_map;
-
     P* channel_state;
     size_t current_channel_id{UNSET_CHANNEL};
+    size_t current_address_space{};
     Tegra::Engines::Maxwell3D* maxwell3d;
     Tegra::Engines::KeplerCompute* kepler_compute;
     Tegra::MemoryManager* gpu_memory;
+
+    std::deque<P> channel_storage;
+    std::deque<size_t> free_channel_ids;
+    std::unordered_map<s32, size_t> channel_map;
+    struct AddresSpaceRef {
+        size_t ref_count;
+        size_t storage_id;
+        Tegra::MemoryManager* gpu_memory;
+    };
+    std::unordered_map<size_t, AddresSpaceRef> address_spaces;
+    mutable std::mutex config_mutex;
+
+    virtual void OnGPUASRegister([[maybe_unused]] size_t map_id) {}
 };
 
 } // namespace VideoCommon